import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import chirp, gauss_spline
class SignalBuilder:
    """
    信号构建器，支持添加多种特征并生成最终信号
    """
    def __init__(self, sampling_rate, duration):
        self.sampling_rate = sampling_rate
        self.duration = duration
        self.num_samples = int(sampling_rate * duration)
        self.time = np.linspace(0, duration, self.num_samples, endpoint=False)
        self.signal = np.zeros(self.num_samples)
        self.features = []

    def add_feature(self, feature_type, **kwargs):
        feature_func = self._feature_factory(feature_type)
        if feature_func:
            self.signal += feature_func(**kwargs)
            self.features.append((feature_type, kwargs))

    def _feature_factory(self, feature_type):
        """工厂方法，返回特征生成函数"""
        return {
            'baseline': self._baseline,
            'noise': self._noise,
            'spike': self._spike,
            'oscillation': self._oscillation
        }.get(feature_type, None)

    def _baseline(self, amplitude=0.5, freq=0.1, offset=5):
        return offset + amplitude * np.sin(2 * np.pi * freq * self.time)

    def _noise(self, std=0.5):
        return std * np.random.randn(self.num_samples)

    def _spike(self, start_time=1.5, duration=0.05, amplitude=10):
        spike = np.zeros(self.num_samples)
        center = int(start_time * self.sampling_rate + duration * self.sampling_rate / 2)
        width = int(duration * self.sampling_rate / 4)
        if 0 <= center < self.num_samples:
            t = np.linspace(-duration/2, duration/2, width * 2)
            pulse = amplitude * (np.exp(-(t*15)**2) - 0.5 * np.exp(-(t*5)**2))
            start_idx = max(center - width, 0)
            end_idx = min(center + width, self.num_samples)
            pulse_len = end_idx - start_idx
            if pulse.shape[0] != pulse_len:
                pulse = np.interp(np.linspace(0,1,pulse_len), np.linspace(0,1,pulse.shape[0]), pulse)
            spike[start_idx:end_idx] = pulse
        return spike

    def _oscillation(self, start_time=3.0, duration=0.8, freq=10, amplitude=8):
        osc = np.zeros(self.num_samples)
        center = int(start_time * self.sampling_rate + duration * self.sampling_rate / 2)
        half_width = int(duration * self.sampling_rate / 2)
        t = np.linspace(-duration/2, duration/2, half_width * 2)
        envelope = np.exp(-(t*4)**2)
        wave = amplitude * np.sin(2 * np.pi * freq * t) * envelope
        start_idx = max(center - half_width, 0)
        end_idx = min(center + half_width, self.num_samples)
        wave_len = end_idx - start_idx
        if wave.shape[0] != wave_len:
            wave = np.interp(np.linspace(0,1,wave_len), np.linspace(0,1,wave.shape[0]), wave)
        osc[start_idx:end_idx] = wave
        return osc

    def get_signal(self):
        return self.time, self.signal

def plot_signal(time, signal, spike_time, spike_duration, osc_time, osc_duration):
    """
    绘制信号及特征区域
    """
    plt.figure(figsize=(15, 6))
    plt.plot(time, signal, lw=1)
    plt.title('Simulated Signal with Different Features')
    plt.xlabel('Time (s)')
    plt.ylabel('Amplitude')
    plt.grid(True)
    plt.axvspan(spike_time, spike_time + spike_duration, color='red', alpha=0.2, label='Spike Region')
    plt.axvspan(osc_time, osc_time + osc_duration, color='blue', alpha=0.2, label='Oscillation Region')
    plt.legend()
    plt.tight_layout()
    plt.show()

# 主流程
if __name__ == '__main__':
    sampling_rate = 1000
    duration = 5
    builder = SignalBuilder(sampling_rate, duration)
    builder.add_feature('baseline', amplitude=0.0, freq=0.1, offset=5)
    builder.add_feature('noise', std=0.5)
    builder.add_feature('spike', start_time=1.5, duration=0.05, amplitude=10)
    builder.add_feature('oscillation', start_time=3.0, duration=0.8, freq=10, amplitude=8)
    time, combined_signal = builder.get_signal()

    plot_signal(time, combined_signal, 1.5, 0.05, 3.0, 0.8)


# 你也可以保存这个信号到文件，例如CSV
# np.savetxt('simulated_signal.csv', combined_signal, delimiter=',')
# print("Signal saved to simulated_signal.csv")